from datetime import datetime, timedelta
from typing import Optional
from fastapi import Request, HTTPException, status
from starlette.middleware.base import BaseHTTPMiddleware
import redis.asyncio as redis
from src.core.config import settings
from src.core.dependencies import get_redis_client
import logging

logger = logging.getLogger(__name__)


class RateLimiter:
    """Rate limiter using Redis."""
    
    def __init__(self, redis_client: redis.Redis, max_requests: int = 100, window_seconds: int = 60):
        self.redis_client = redis_client
        self.max_requests = max_requests
        self.window_seconds = window_seconds
    
    async def check_rate_limit(self, identifier: str) -> tuple[bool, int]:
        """
        Check if the identifier has exceeded the rate limit.
        Returns (is_allowed, remaining_requests).
        """
        key = f"rate_limit:{identifier}"
        
        try:
            # Use pipeline for atomic operations
            async with self.redis_client.pipeline() as pipe:
                # Get current count
                pipe.get(key)
                # Increment count
                pipe.incr(key)
                # Set expiry if new key
                pipe.expire(key, self.window_seconds)
                
                results = await pipe.execute()
                
                current_count = int(results[1])  # After increment
                
                if current_count == 1:
                    # New key, set expiry
                    await self.redis_client.expire(key, self.window_seconds)
                
                if current_count > self.max_requests:
                    return False, 0
                
                return True, self.max_requests - current_count
                
        except Exception as e:
            logger.error(f"Rate limiting error: {str(e)}")
            # On error, allow the request
            return True, self.max_requests
    
    async def reset_limit(self, identifier: str) -> None:
        """Reset rate limit for an identifier."""
        key = f"rate_limit:{identifier}"
        await self.redis_client.delete(key)


class RateLimitMiddleware(BaseHTTPMiddleware):
    """Middleware for API rate limiting."""
    
    async def dispatch(self, request: Request, call_next):
        # Skip rate limiting for health checks
        if request.url.path in ["/health", "/metrics"]:
            return await call_next(request)
        
        # Get client identifier (IP address)
        client_ip = request.client.host if request.client else "unknown"
        
        # Get redis client
        redis_client = await get_redis_client()
        
        # Create rate limiter
        rate_limiter = RateLimiter(
            redis_client,
            max_requests=settings.RATE_LIMIT_PER_MINUTE,
            window_seconds=60
        )
        
        # Check rate limit
        is_allowed, remaining = await rate_limiter.check_rate_limit(client_ip)
        
        if not is_allowed:
            logger.warning(f"Rate limit exceeded for IP: {client_ip}")
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail="Rate limit exceeded. Please try again later.",
                headers={
                    "X-RateLimit-Limit": str(settings.RATE_LIMIT_PER_MINUTE),
                    "X-RateLimit-Remaining": "0",
                    "X-RateLimit-Reset": str(int(datetime.utcnow().timestamp()) + 60)
                }
            )
        
        # Process request
        response = await call_next(request)
        
        # Add rate limit headers
        response.headers["X-RateLimit-Limit"] = str(settings.RATE_LIMIT_PER_MINUTE)
        response.headers["X-RateLimit-Remaining"] = str(remaining)
        response.headers["X-RateLimit-Reset"] = str(int(datetime.utcnow().timestamp()) + 60)
        
        return response


async def check_login_attempts(identifier: str, redis_client: redis.Redis) -> bool:
    """Check if login attempts exceeded."""
    attempts_key = f"login_attempts:{identifier}"
    attempts = await redis_client.get(attempts_key)
    
    if attempts and int(attempts) >= settings.LOGIN_ATTEMPT_LIMIT:
        return False
    return True


async def increment_login_attempts(identifier: str, redis_client: redis.Redis) -> None:
    """Increment failed login attempts."""
    attempts_key = f"login_attempts:{identifier}"
    
    async with redis_client.pipeline() as pipe:
        pipe.incr(attempts_key)
        pipe.expire(attempts_key, settings.LOGIN_ATTEMPT_WINDOW_MINUTES * 60)
        await pipe.execute()


async def reset_login_attempts(identifier: str, redis_client: redis.Redis) -> None:
    """Reset login attempts on successful login."""
    attempts_key = f"login_attempts:{identifier}"
    await redis_client.delete(attempts_key)